from PyQt5.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QComboBox,
    QTableWidget, QTableWidgetItem, QHeaderView, QAbstractItemView, QFileDialog,
    QMessageBox
)


class SubtitleEditorWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("字幕编辑器（SRT/ASS）")
        self.resize(900, 520)
        self._build_ui()
        self._wire()

    def _build_ui(self):
        self.format_combo = QComboBox(self)
        self.format_combo.addItems(["SRT", "ASS"])
        self.open_btn = QPushButton("打开", self)
        self.save_btn = QPushButton("保存", self)
        self.add_row_btn = QPushButton("添加一行", self)
        self.del_row_btn = QPushButton("删除选中", self)

        top = QHBoxLayout()
        top.addWidget(QLabel("格式:"))
        top.addWidget(self.format_combo)
        top.addStretch(1)
        top.addWidget(self.add_row_btn)
        top.addWidget(self.del_row_btn)
        top.addWidget(self.open_btn)
        top.addWidget(self.save_btn)

        self.table = QTableWidget(self)
        self.table.setColumnCount(3)
        self.table.setHorizontalHeaderLabels(["开始(00:00:00,000)", "结束(00:00:00,000)", "文本"])
        header = self.table.horizontalHeader()
        header.setSectionResizeMode(0, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(1, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(2, QHeaderView.Stretch)
        self.table.setEditTriggers(QAbstractItemView.AllEditTriggers)

        layout = QVBoxLayout(self)
        layout.addLayout(top)
        layout.addWidget(self.table)
        self.setLayout(layout)

    def _wire(self):
        self.add_row_btn.clicked.connect(self._add_row)
        self.del_row_btn.clicked.connect(self._delete_selected)
        self.open_btn.clicked.connect(self._open_file)
        self.save_btn.clicked.connect(self._save_file)

    def _add_row(self):
        r = self.table.rowCount()
        self.table.insertRow(r)
        self.table.setItem(r, 0, QTableWidgetItem("00:00:00,000"))
        self.table.setItem(r, 1, QTableWidgetItem("00:00:03,000"))
        self.table.setItem(r, 2, QTableWidgetItem("字幕内容"))

    def _delete_selected(self):
        rows = sorted({i.row() for i in self.table.selectedIndexes()}, reverse=True)
        for r in rows:
            self.table.removeRow(r)

    @staticmethod
    def _read_text(path: str) -> str:
        with open(path, 'r', encoding='utf-8', errors='ignore') as f:
            return f.read()

    def _open_file(self):
        path, _ = QFileDialog.getOpenFileName(self, "打开字幕", "", "字幕文件 (*.srt *.ass)")
        if not path:
            return
        content = self._read_text(path)
        if path.lower().endswith('.srt'):
            self.format_combo.setCurrentText('SRT')
            self._load_srt(content)
        else:
            self.format_combo.setCurrentText('ASS')
            self._load_ass(content)

    def _save_file(self):
        fmt = self.format_combo.currentText().upper()
        if fmt == 'SRT':
            path, _ = QFileDialog.getSaveFileName(self, "保存 SRT", "", "SRT (*.srt)")
            if not path:
                return
            text = self._to_srt()
        else:
            path, _ = QFileDialog.getSaveFileName(self, "保存 ASS", "", "ASS (*.ass)")
            if not path:
                return
            text = self._to_ass()
        with open(path, 'w', encoding='utf-8') as f:
            f.write(text)
        QMessageBox.information(self, "成功", f"已保存到:\n{path}")

    def _load_srt(self, content: str):
        self.table.setRowCount(0)
        blocks = [b for b in content.split('\n\n') if b.strip()]
        for b in blocks:
            lines = [l for l in b.splitlines() if l.strip()]
            if len(lines) < 2:
                continue
            if '-->' in lines[0]:
                time_line = lines[0]
                text_lines = lines[1:]
            else:
                time_line = lines[1] if len(lines) > 1 else ''
                text_lines = lines[2:]
            if '-->' not in time_line:
                continue
            start, end = [t.strip() for t in time_line.split('-->')]
            text = '\n'.join(text_lines)
            r = self.table.rowCount()
            self.table.insertRow(r)
            self.table.setItem(r, 0, QTableWidgetItem(start))
            self.table.setItem(r, 1, QTableWidgetItem(end))
            self.table.setItem(r, 2, QTableWidgetItem(text))

    def _to_srt(self) -> str:
        lines = []
        for i in range(self.table.rowCount()):
            start = (self.table.item(i, 0) or QTableWidgetItem("")).text()
            end = (self.table.item(i, 1) or QTableWidgetItem("")).text()
            text = (self.table.item(i, 2) or QTableWidgetItem("")).text()
            lines.append(str(i + 1))
            lines.append(f"{start} --> {end}")
            lines.extend(text.splitlines() or [""])
            lines.append("")
        return '\n'.join(lines).strip() + '\n'

    def _load_ass(self, content: str):
        self.table.setRowCount(0)
        in_events = False
        fields = []
        for line in content.splitlines():
            if line.strip().lower().startswith('[events]'):
                in_events = True
                continue
            if not in_events:
                continue
            if line.strip().lower().startswith('format:'):
                fields = [f.strip().lower() for f in line.split(':', 1)[1].split(',')]
                continue
            if line.strip().lower().startswith('dialogue:'):
                parts = line.split(':', 1)[1].split(',', len(fields) - 1)
                m = dict(zip(fields, parts))
                start = m.get('start', '0:00:00.00').replace('.', ',')
                end = m.get('end', '0:00:03.00').replace('.', ',')
                text = m.get('text', '').replace('\\N', '\n')
                r = self.table.rowCount()
                self.table.insertRow(r)
                self.table.setItem(r, 0, QTableWidgetItem(start))
                self.table.setItem(r, 1, QTableWidgetItem(end))
                self.table.setItem(r, 2, QTableWidgetItem(text))

    def _to_ass(self) -> str:
        header = (
            "[Script Info]\n"
            "ScriptType: v4.00+\n"
            "PlayResX: 1920\n"
            "PlayResY: 1080\n"
            "WrapStyle: 0\n"
            "ScaledBorderAndShadow: yes\n\n"
            "[V4+ Styles]\n"
            "Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding\n"
            "Style: Default,Arial,40,&H00FFFFFF,&H000000FF,&H00000000,&H64000000,0,0,0,0,100,100,0,0,1,2,0,2,10,10,20,1\n\n"
            "[Events]\n"
            "Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n"
        )
        lines = [header]
        for i in range(self.table.rowCount()):
            start = (self.table.item(i, 0) or QTableWidgetItem("")).text().replace(',', '.')
            end = (self.table.item(i, 1) or QTableWidgetItem("")).text().replace(',', '.')
            text = (self.table.item(i, 2) or QTableWidgetItem("")).text().replace('\n', r'\\N')
            lines.append(f"Dialogue: 0,{start},{end},Default,,0,0,20,,{text}")
        return '\n'.join(lines) + '\n'


